#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
:File       :mnist.py
:Description:
    MINIST 数据加载器
    1. torchvision 库中实现了 MNIST数据集的下载以及载入逻辑 [see torchvision.datasets.MNIST]
        但是,该MNIST类只适用于 二进制形式存储的MNIST数据集。
        针对已经可视化的MNIST数据集来说,需要自己来实现具体的载入逻辑,即 mnist.py实现的核心功能

    2. 可视化MNIST数据集的数据组织形式:
        mnist
         ├─ train
         │   └─ <label-index>.jpg
         └─ test
             └─ <label-index>.jpg

:EditTime   :2025/05/03 18:21:46
:Author     :Kiumb
'''
import os
import torch
from typing import *
import torchvision.transforms.functional as T
import torchvision.io.image as I 

class MNIST(torch.utils.data.Dataset):
    def __init__(self,root:str,type:str):
        '''
        root 为 MNIST数据集存放的根目录
        type 指定 加载MNIST数据集的训练集部分还是测试集部分
        '''
        assert type in ['train','test'] 
        
        self.img_list = []
        for img_name in os.listdir(root + os.sep + type):
            if img_name.endswith('.jpg'):
                self.img_list.append(root + os.sep + type + os.sep + img_name)  

    def __len__(self):
        return len(self.img_list)

    def __getitem__(self,idx:int) -> Tuple[torch.Tensor,torch.Tensor]:    
        # No data augmentation is applied to the MNIST dataset

        img_path = self.img_list[idx]

        img_tensor = I.read_image(img_path).to(torch.float32) / 255  # The size of img_tensor is [1,28,28]
        # further normalized the img tensor 
        img_tensor = T.normalize(img_tensor,mean=(0.1307,),std=(0.3081,),inplace=True )

        # GET LABEL
        label = torch.as_tensor(int(img_path.split(os.sep)[-1][0])) 
    
        return img_tensor,label


